/*
 *
 *  ------------------------------------------------------------------
 *  Copyright © 2017 Hangzhou DtDream Technology Co.,Lt d. All rights reserved.
 *  ------------------------------------------------------------------
 *        Product:  EMR
 *    Module Name: NEZHA
 *   Date Created: 17-10-23 下午4:22
 *    Description:
 *  ------------------------------------------------------------------
 *  Modification History
 *  DATE            Name           Description
 *  ------------------------------------------------------------------
 *  2017-10-23      NEZHA EMR
 *  ------------------------------------------------------------------
 * /
 */

package org.xukai.remoting.sdk.web;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.sql.SQLException;
import java.util.Arrays;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;
import javax.print.attribute.standard.JobState;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Created by moriarty on 17-2-23.
 */
class SQLClientImpl implements SQLClient {

    private static final Logger LOG = LoggerFactory.getLogger(SQLClientImpl.class);

    private AccessKey accessKey;
    private String database;
    private SQLClientManager sqlClientManager;

    SQLClientImpl(AccessKey accessKey, String database, SQLClientManager sqlClientManager) throws SQLException {
        this.accessKey = accessKey;
        this.database = database;
        this.sqlClientManager = sqlClientManager;
    }

    @Override
    public void execute(String sql) throws SQLException {
        waitForComplete(submitJob(sql));
    }

    private SQLJobReport waitForComplete(String jobId) throws SQLException {
        SQLJobReport jobReport = progress(jobId);
        try {
            while (!JobState.isFinal(jobReport.getState())) {
                LOG.debug("JobReport: JobId={}, JobState={}", jobId, jobReport.getState());
                TimeUnit.SECONDS.sleep(1);
                jobReport = progress(jobId);
            }

            if (JobState.isFailed(jobReport.getState())) {
                SQLJobStmtReport failedStmt = jobReport.getFailedStmt();
                if (failedStmt == null || !StmtState.isFinal(failedStmt.getState())) {
                    throw new SQLException(jobReport.getException());
                } else {
                    throw new SQLException(failedStmt.getException());
                }
            } else if (JobState.isCancelled(jobReport.getState())) {
                throw new SQLException("Statement was canceled");
            }

            return jobReport;
        } catch (InterruptedException e) {
            throw new SQLException("Interrupted while waiting SQLJob", e);
        }
    }

    @Override
    public SQLResult executeQuery(String sql) throws SQLException {
        SQLJobReport jobReport = waitForComplete(submitJob(sql));
        return fetchSQLResult(jobReport.getJobId(), 0);
    }

    @Override
    public List<SQLResult> executeBatch(List<String> sqls) throws SQLException {
        List<SQLResult> result = new LinkedList<>();
        SQLJobReport jobReport = waitForComplete(submitJob(sqls));
        for (int index = 0; index < jobReport.totalStmts(); index++) {
            result.add(fetchSQLResult(jobReport.getJobId(), index));
        }
        return result;
    }

    @Override
    public void addFile(String path) throws SQLException {
        execute("ADD FILE " + path);
    }

    @Override
    public void addJar(String path) throws SQLException {
        execute("ADD JAR " + path);
    }

    @Override
    public void deleteJar(String path) throws SQLException {
        execute("DELETE JAR " + path);
    }

    @Override
    public void createTable(String tableName, String schemaStr) throws SQLException {
        execute("CREATE TABLE " + TableSchema.withinBackticks(tableName) + " (" + schemaStr + ")");
    }

    @Override
    public void createTable(String tableName, TableSchema schema) throws SQLException {
        execute("CREATE TABLE " + TableSchema.withinBackticks(tableName) + schema.getSchemaString());
    }

    @Override
    public void dropTable(String tableName) throws SQLException {
        execute("DROP TABLE " + TableSchema.withinBackticks(tableName));
    }

    @Override
    public void truncateTable(String tableName) throws SQLException {
        execute("TRUNCATE TABLE " + TableSchema.withinBackticks(tableName));
    }

    @Override
    public void addFunction(String funcName, String className) throws SQLException {
        assertNotBuiltin(funcName);
        execute(String.format("CREATE FUNCTION %s AS '%s'", funcName, className));
    }

    @Override
    public void addFunction(String funcName, String className, String jarPath) throws SQLException {
        assertNotBuiltin(funcName);
        execute(String.format("CREATE FUNCTION %s AS '%s' USING JAR '%s'", funcName, className,
                jarPath));
    }

    @Override
    public void dropFunction(String funcName) throws SQLException {
        execute(String.format("DROP FUNCTION IF EXISTS %s", funcName));
    }

    @Override
    public SQLResult showPartitions(String tableName) throws SQLException {
        return showPartitions(tableName, null);
    }

    @Override
    public SQLResult showPartitions(String tableName, Map<String, String> partSpec) throws SQLException {
        try {
            return sqlClientManager.restClient.showPartitions(accessKey, database, tableName, partSpec);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public void createPartition(String tableName, Map<String, String> partName) throws SQLException {
        try {
            sqlClientManager.restClient.createPartition(accessKey, database, tableName, partName);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public void dropPartition(String tableName, Map<String, String> partition) throws SQLException {
        execute(String.format("ALTER TABLE %s DROP PARTITION %s",
                TableSchema.withinBackticks(tableName),
                TableSchema.partSpec(partition)));
    }

    @Override
    public void truncatePartition(String tableName, Map<String, String> partition) throws SQLException {
        execute(String.format("TRUNCATE TABLE %s PARTITION %s",
                TableSchema.withinBackticks(tableName),
                TableSchema.partSpec(partition)));
    }

    @Override
    public List<String> getAllTableNames() throws SQLException {
        try {
            return sqlClientManager.restClient.getTableNames(accessKey, database, null, new String[0]);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public List<String> getTableTypes() throws SQLException {
        try {
            return sqlClientManager.restClient.getTableTypes(accessKey, database);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public List<String> getTableNames(String tablePattern, String... tableType) throws SQLException {
        try {
            return sqlClientManager.restClient.getTableNames(accessKey, database, tablePattern, tableType);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    @Deprecated
    public List<TableDesc> getTables() throws SQLException {
        try {
            return sqlClientManager.restClient.getTables(accessKey, database, null, new String[0]);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    @Deprecated
    public List<TableDesc> getTables(String tablePattern, String... tableType) throws SQLException {
        try {
            return sqlClientManager.restClient.getTables(accessKey, database, tablePattern, tableType);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public Page<TableDesc> getTables(int pageNum, int pageSize, String tablePattern, String... tableType)
            throws SQLException {
        try {
            return sqlClientManager.restClient.getTablesPaged(accessKey, database, tablePattern, tableType, pageNum,
                    pageSize);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public TableDetail getTableDetail(String tableName) throws SQLException {
        try {
            return sqlClientManager.restClient.getTableDetail(accessKey, database, tableName);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public DatabaseDetail getDatabaseDetail(String databaseName) throws SQLException {
        try {
            return sqlClientManager.restClient.getDatabaseDetail(accessKey, databaseName);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public SQLResult getTableRecord(String tableName, String[] columns) throws SQLException {
        String cols = join(columns, ',');
        return executeQuery(String.format("SELECT %s FROM ", cols) + TableSchema.withinBackticks(tableName));
    }

    @Override
    public SQLResult getTableRecord(String tableName, String[] columns, int limit) throws SQLException {
        String cols = join(columns, ',');
        return executeQuery(String.format("SELECT %s FROM ", cols) + TableSchema.withinBackticks(tableName) +
                " LIMIT " + limit);
    }

    private static String join(String[] columns, char separator) {
        if (columns == null) {
            return null;
        }

        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < columns.length; i++) {
            if (i > 0) {
                sb.append(separator);
            }

            if (columns[i] != null) {
                sb.append(columns[i]);
            }
        }
        return sb.toString();
    }

    @Override
    @Deprecated
    public SQLResult getTableRecord(String tableName) throws SQLException {
        return executeQuery("SELECT * FROM " + TableSchema.withinBackticks(tableName));
    }

    @Override
    @Deprecated
    public SQLResult getTableRecord(String tableName, int limit) throws SQLException {
        return executeQuery("SELECT * FROM " + TableSchema.withinBackticks(tableName) + " LIMIT " + limit);
    }

    @Override
    public void refreshTable(String tableName) throws SQLException {
        execute("REFRESH TABLE " + TableSchema.withinBackticks(tableName));
    }

    @Override
    public String submitJob(String sql) throws SQLException {
        return submitJob(Arrays.asList(sql));
    }

    @Override
    public String submitJob(List<String> sqls) throws SQLException {
        try {
            return sqlClientManager.restClient.submitJob(accessKey, database, sqls);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public SQLJobReport progress(String jobId) throws SQLException {
        try {
            return sqlClientManager.restClient.getJobReport(accessKey, database, jobId);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public SQLResult fetchSQLResult(String jobId, int stmtIndex) throws SQLException {
        try {
            SQLJobStmtResult stmt = sqlClientManager.restClient.getJobStmtResult(accessKey, database, jobId, stmtIndex);
            switch (stmt.getState()) {
                case INITIALIZED:
                case RUNNING:
                case PENDING:
                    throw new SQLException("Statement has not finished.");
                case CANCELLED:
                    throw new SQLException("Statement was canceled");
                case ERROR:
                    throw new SQLException(stmt.getException());
                case FINISHED:
                    return stmt.getSqlResult();
                default:
                    throw new SQLException("Unknown Statement state " + stmt.getState());
            }

        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public SQLLog fetchSQLLog(String jobId, int stmtIndex) throws SQLException {
        try {
            SQLJobStmtLog stmtLog = sqlClientManager.restClient.getJobStmtLog(accessKey, database, jobId, stmtIndex);
            return new SQLLog(stmtLog);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public SQLTracker fetchSQLTracker(String jobId, int stmtIndex) throws SQLException {
        return new SQLTracker(accessKey, sqlClientManager.restClient, database, jobId, stmtIndex);
    }

    @Override
    public void cancelJob(String jobId) throws SQLException {
        try {
            sqlClientManager.restClient.cancelJob(accessKey, database, jobId);
        } catch (NeZhaClientException e) {
            throw new SQLException(e.getMessage(), e);
        }
    }

    @Override
    public void close() throws SQLException {
    }

    public static void assertNotBuiltin(String funcName) throws SQLException {
        if (funcName != null && BUILTIN_FUNCTIONS.contains(funcName.trim().toLowerCase())) {
            throw new SQLException("Function name conflicts with builtin function " + funcName.trim().toLowerCase());
        }
    }

    // Builtin functions should keep consistent with DTSpark
    public static final LinkedHashSet<String> BUILTIN_FUNCTIONS = new LinkedHashSet<>();

    static {
        try (BufferedReader bufReader = new BufferedReader(
                new InputStreamReader(SQLClientImpl.class.getResourceAsStream("/builtins.txt"), StandardCharsets.UTF_8))) {
            for (String line = ""; line != null; line = bufReader.readLine()) {
                line = line.trim();
                if (!line.isEmpty()) {
                    BUILTIN_FUNCTIONS.add(line.toLowerCase());
                }
            }
        } catch (IOException e) {
            throw new Error(e);
        }
    }
}
